import torch
from torch import nn
from .. import base
import torch.nn.functional as F


def get_synapse_neuron_list_from_model(model: nn.Module) -> list:
    snlist = []
    modulelist = list(model.modules())
    
    for i in range(len(modulelist)):
        if isinstance(modulelist[i], (nn.Linear, nn.Conv2d)):
            snlist.append({'synapse' : modulelist[i], 'neuron' : modulelist[i + 1]})
    return snlist

class OnlineSupervisedDtrpSTDPLearner(nn.Module):
    def __init__(self, model: nn.Module, lr: float, decay: float, Bthres: float = 0.7):
        super().__init__()
        self.lr = lr
        self.trace_decay = (decay - 1) / decay
        self.Bthres = Bthres
        self.snlist = get_synapse_neuron_list_from_model(model)
        for item in self.snlist:
            item['ri'] = 0.
            item['rj'] = 0.
            item['hat_rj'] = 0.
            item['synapse'].forward = self.input_monitor(item['synapse'].forward, item)
            item['neuron'].forward = self.output_monitor(item['neuron'].forward, item)
            item['Bmem'] = 0.

    def getB(self):
        for item in self.snlist:
            if isinstance(item['synapse'], nn.Linear):
                item['B'] = torch.rand(size=(self.snlist[-1]['sj'].shape[1],item['sj'].shape[1]))
            elif isinstance(item['synapse'], nn.Conv2d):
                item['B'] = torch.rand(size=(self.snlist[-1]['sj'].shape[1], item['sj'].shape[1]*item['sj'].shape[2]*item['sj'].shape[3]))
            item['B'] = torch.where(item['B'] >= self.Bthres, item['B'], torch.zeros_like(item['B'])).to(item['synapse'].weight.device)


    def input_monitor(self, func, sn):
        def wrapper(*args, **kwargs):
            sn['si'] = args[0]
            result = func(*args, **kwargs)
            return result
        return wrapper
    
    def output_monitor(self, func, sn):
        def wrapper(*args, **kwargs):
            result = func(*args, **kwargs)
            sn['sj'] = result
            return result
        return wrapper
    
    def linear_dtrp_stdp(self, item: dict):
        sj = item['sj']
        si = item['si']
        hat_sj = item['hat_sj']
        ri = item['ri'] if not isinstance(item['ri'], float) else torch.zeros_like(si)
        rj = item['rj'] if not isinstance(item['rj'], float) else torch.zeros_like(sj)
        hat_rj = item['hat_rj'] if not isinstance(item['hat_rj'], float) else torch.zeros_like(sj)
        delta_s = hat_sj - sj
        delta_r = hat_rj - rj
        update1 = (delta_s.unsqueeze(2) * ri.unsqueeze(1)).sum(dim=0)
        update2 = (delta_r.unsqueeze(2) * si.unsqueeze(1)).sum(dim=0)
        item['synapse'].weight.data.add_((update1 + update2), alpha=self.lr)

        item['ri'] = item['ri'] * self.trace_decay + si
        item['rj'] = item['rj'] * self.trace_decay + sj
        item['hat_rj'] = item['hat_rj'] * self.trace_decay + hat_sj

    def conv2d_drtp_stdp(self, item: dict):
        """测试版本，最好只用stride=1,padding=1,kernel_size=3的卷积层"""
        stride_h, stride_w = item['synapse'].stride
        si = item['si']
        if item['synapse'].padding != (0, 0):
            ph, pw = item['synapse'].padding
            if item['synapse'].padding_mode != 'zeros':
                raise NotImplementedError('不支持非0填充')
            else:
                si = F.pad(item['si'], pad=(pw, pw, ph, ph))
        sj = item['sj']
        hat_sj = item['hat_sj']
        ri = item['ri'] if not isinstance(item['ri'], float) else torch.zeros_like(si)
        rj = item['rj'] if not isinstance(item['rj'], float) else torch.zeros_like(sj)
        hat_rj = item['hat_rj'] if not isinstance(item['hat_rj'], float) else torch.zeros_like(sj)
        
        delta_w = torch.zeros_like(item['synapse'].weight.data) 
        for h in range(item['synapse'].weight.shape[2]):
            for w in range(item['synapse'].weight.shape[3]):
                h_end = si.shape[2] - item['synapse'].weight.shape[2] + 1 + h
                w_end = si.shape[3] - item['synapse'].weight.shape[3] + 1 + w
                sihw = si[:, :, h:h_end:stride_h, w:w_end:stride_w]
                rihw = ri[:, :, h:h_end:stride_h, w:w_end:stride_w]
                delta_s = hat_sj - sj
                delta_r = hat_rj - rj
                update1 = (delta_s.unsqueeze(2) * rihw.unsqueeze(1)).permute(1, 2, 0, 3, 4).sum(dim=[2, 3, 4])
                update2 = (delta_r.unsqueeze(2) * sihw.unsqueeze(1)).permute(1, 2, 0, 3, 4).sum(dim=[2, 3, 4])
                delta_w[:, :, h, w] += update1 + update2
        item['synapse'].weight.data.add_(delta_w, alpha=self.lr)


        item['ri'] = item['ri'] * self.trace_decay + si
        item['rj'] = item['rj'] * self.trace_decay + sj
        item['hat_rj'] = item['hat_rj'] * self.trace_decay + hat_sj


    def conv2d_stdp(self, item: dict):
        """测试版本，最好只用stride=1,padding=1,kernel_size=3的卷积层"""
        stride_h, stride_w = item['synapse'].stride
        si = item['si']
        if item['synapse'].padding != (0, 0):
            ph, pw = item['synapse'].padding
            if item['synapse'].padding_mode != 'zeros':
                raise NotImplementedError('不支持非0填充')
            else:
                si = F.pad(item['si'], pad=(pw, pw, ph, ph))
        sj = item['sj']
        ri = item['ri'] if not isinstance(item['ri'], float) else torch.zeros_like(si)
        rj = item['rj'] if not isinstance(item['rj'], float) else torch.zeros_like(sj)
        
        delta_w = torch.zeros_like(item['synapse'].weight.data) 
        for h in range(item['synapse'].weight.shape[2]):
            for w in range(item['synapse'].weight.shape[3]):
                h_end = si.shape[2] - item['synapse'].weight.shape[2] + 1 + h
                w_end = si.shape[3] - item['synapse'].weight.shape[3] + 1 + w
                sihw = si[:, :, h:h_end:stride_h, w:w_end:stride_w]
                rihw = ri[:, :, h:h_end:stride_h, w:w_end:stride_w]
                update1 = (rj.unsqueeze(2) * sihw.unsqueeze(1)).permute(1, 2, 0, 3, 4).sum(dim=[2, 3, 4])
                update2 = (sj.unsqueeze(2) * rihw.unsqueeze(1)).permute(1, 2, 0, 3, 4).sum(dim=[2, 3, 4])
                delta_w[:, :, h, w] += update1 - update2
        item['synapse'].weight.data.add_(delta_w, alpha=self.lr)


        item['ri'] = item['ri'] * self.trace_decay + si
        item['rj'] = item['rj'] * self.trace_decay + sj
    
        
    def step(self, so, hat_so):
        #如果是首次调整，需要先决定一下随机矩阵B
        if not 'B' in self.snlist[-1]:
            self.getB()

        for i, item in enumerate(self.snlist):
            if i == len(self.snlist) - 1:
                item['hat_sj'] = hat_so
                self.linear_dtrp_stdp(item)
            else:
                if isinstance(item['synapse'], nn.Linear):
                    item['Bmem'] = torch.bernoulli(hat_so @ item['B']) + item['Bmem']
                    item['hat_sj'] = torch.zeros_like(item['Bmem'])
                    item['hat_sj'][item['Bmem']>=1.] = 1
                    item['Bmem'][item['Bmem']>=1.] -= 1
                    self.linear_dtrp_stdp(item)
                # elif isinstance(item['synapse'], nn.Conv2d):
                #     item['hat_sj'] = (hat_so @ item['B']).view(item['sj'].shape)
                #     item['hat_sj'] = torch.bernoulli(item['hat_sj'])
                #     self.conv2d_stdp(item)
            # print(f"item {i}: sj={sj.shape} hat_sj={hat_sj.shape}")
    
    def reset(self):
        for item in self.snlist:
            item['ri'] = 0.
            item['rj'] = 0.
            item['hat_rj'] = 0.
            item['Bmem'] = 0.